"""
计算引擎模块，负责执行测算逻辑
"""
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional
import logging
from datetime import datetime

def _handle_time_period(self, start_date: str, periods: int, freq: str) -> pd.DatetimeIndex:
        """统一处理时间周期生成"""
        return pd.date_range(
            start=pd.to_datetime(start_date) + pd.DateOffset(months=1),
            periods=periods,
            freq=freq
        )

def _format_result(self, data: Any, metadata: Dict[str, Any]) -> Dict[str, Any]:
        """统一格式化计算结果"""
        return {
            "data": data,
            "metadata": metadata,
            "timestamp": datetime.now().isoformat()
        }

from models.scene import Scene
from models.calculation import CalculationResult

logger = logging.getLogger(__name__)


class CalculationEngine:
    """
    计算引擎，负责执行测算逻辑并生成结果
    """
    
    def __init__(self):
        """
        初始化计算引擎
        """
        # 注册可用的计算方法
        self.calculation_methods = {
            "sales_forecast": self._calculate_sales_forecast,
            "investment_return": self._calculate_investment_return,
            "cost_analysis": self._calculate_cost_analysis,
            "profit_projection": self._calculate_profit_projection,
            "custom": self._calculate_custom
        }
    
    def execute_calculation(self, scene: Scene, data: Dict[str, pd.DataFrame]) -> CalculationResult:
        """
        执行测算
        
        Args:
            scene: 场景对象
            data: 输入数据，字典形式，键为数据源ID，值为DataFrame
            
        Returns:
            计算结果对象
        """
        start_time = datetime.now()
        
        try:
            # 根据场景类型选择计算方法
            if scene.scene_type in self.calculation_methods:
                calculation_method = self.calculation_methods[scene.scene_type]
                result_data = calculation_method(scene, data)
            else:
                raise ValueError(f"不支持的场景类型: {scene.scene_type}")
            
            # 创建计算结果
            end_time = datetime.now()
            execution_time = (end_time - start_time).total_seconds()
            
            result = CalculationResult(
                id=f"calc_{scene.id}_{start_time.strftime('%Y%m%d%H%M%S')}",
                scene_id=scene.id,
                start_time=start_time.isoformat(),
                end_time=end_time.isoformat(),
                execution_time=execution_time,
                status="success",
                result_data=result_data,
                error=None
            )
            
            return result
            
        except Exception as e:
            # 记录错误
            logger.error(f"计算错误: {str(e)}", exc_info=True)
            
            # 创建错误结果
            end_time = datetime.now()
            execution_time = (end_time - start_time).total_seconds()
            
            result = CalculationResult(
                id=f"calc_{scene.id}_{start_time.strftime('%Y%m%d%H%M%S')}",
                scene_id=scene.id,
                start_time=start_time.isoformat(),
                end_time=end_time.isoformat(),
                execution_time=execution_time,
                status="error",
                result_data={},
                error=str(e)
            )
            
            return result
    
    def _calculate_sales_forecast(self, scene: Scene, data: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
        """
        销售预测计算逻辑
        
        Args:
            scene: 场景对象
            data: 输入数据
            
        Returns:
            计算结果数据
        """
        # 获取参数
        params = scene.parameters
        growth_rate = params.get("growth_rate", 0.05)
        forecast_periods = params.get("forecast_periods", 12)
        seasonality = params.get("seasonality", False)
        
        # 获取历史销售数据
        if "sales_history" not in data:
            raise ValueError("缺少历史销售数据")
        
        sales_history = data["sales_history"]
        
        # 简单的销售预测逻辑
        last_period_sales = sales_history["sales"].iloc[-1]
        
        # 创建预测结果
        forecast_dates = pd.date_range(
            start=pd.to_datetime(sales_history.index[-1]) + pd.DateOffset(months=1),
            periods=forecast_periods,
            freq='M'
        )
        
        forecast_sales = []
        for i in range(forecast_periods):
            # 考虑季节性因素
            if seasonality and len(sales_history) >= 12:
                # 使用去年同期数据作为基础
                season_index = i % 12
                base_sales = sales_history["sales"].iloc[-(12 - season_index)]
                forecasted_value = base_sales * (1 + growth_rate) ** (i // 12 + 1)
            else:
                # 简单的指数增长
                forecasted_value = last_period_sales * (1 + growth_rate) ** (i + 1)
            
            forecast_sales.append(forecasted_value)
        
        # 创建预测DataFrame
        forecast_df = pd.DataFrame({
            "date": forecast_dates,
            "sales": forecast_sales
        })
        
        # 计算汇总统计
        total_forecast = sum(forecast_sales)
        average_forecast = total_forecast / forecast_periods
        max_forecast = max(forecast_sales)
        min_forecast = min(forecast_sales)
        
        # 返回结果
        return {
            "forecast_data": forecast_df.to_dict(orient="records"),
            "summary": {
                "total_forecast": total_forecast,
                "average_forecast": average_forecast,
                "max_forecast": max_forecast,
                "min_forecast": min_forecast,
                "growth_rate": growth_rate,
                "forecast_periods": forecast_periods
            },
            "charts": {
                "historical_and_forecast": {
                    "type": "line",
                    "data": {
                        "historical": sales_history["sales"].tolist(),
                        "forecast": forecast_sales
                    },
                    "labels": [d.strftime("%Y-%m") for d in forecast_dates]
                }
            }
        }
    
    def _calculate_investment_return(self, scene: Scene, data: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
        """
        投资回报分析计算逻辑
        
        Args:
            scene: 场景对象
            data: 输入数据
            
        Returns:
            计算结果数据
        """
        # 获取参数
        params = scene.parameters
        initial_investment = params.get("initial_investment", 100000)
        discount_rate = params.get("discount_rate", 0.08)
        project_years = params.get("project_years", 5)
        
        # 获取现金流数据
        if "cash_flows" in data:
            cash_flows = data["cash_flows"]
            cash_flow_values = cash_flows["cash_flow"].tolist()
        else:
            # 如果没有提供现金流数据，使用参数中的估计值
            estimated_annual_return = params.get("estimated_annual_return", 0.15)
            cash_flow_values = []
            for year in range(1, project_years + 1):
                cash_flow = initial_investment * estimated_annual_return * (1 + 0.05) ** (year - 1)
                cash_flow_values.append(cash_flow)
        
        # 计算净现值(NPV)
        npv = -initial_investment
        for i, cf in enumerate(cash_flow_values):
            npv += cf / ((1 + discount_rate) ** (i + 1))
        
        # 计算内部收益率(IRR)
        # 使用numpy的irr函数
        cash_flows_with_initial = [-initial_investment] + cash_flow_values
        try:
            irr = np.irr(cash_flows_with_initial)
        except:
            irr = None
        
        # 计算回收期
        cumulative_cash_flow = -initial_investment
        payback_period = None
        for i, cf in enumerate(cash_flow_values):
            cumulative_cash_flow += cf
            if cumulative_cash_flow >= 0 and payback_period is None:
                # 简单回收期计算
                payback_period = i + 1
                # 更精确的回收期计算（考虑部分年份）
                if i > 0:
                    previous_cf = cumulative_cash_flow - cf
                    if previous_cf < 0:
                        payback_period = i + abs(previous_cf) / cf
        
        # 计算投资回报率(ROI)
        total_returns = sum(cash_flow_values)
        roi = (total_returns - initial_investment) / initial_investment
        
        # 创建年度现金流表
        yearly_data = []
        cumulative_cf = -initial_investment
        for i, cf in enumerate(cash_flow_values):
            year = i + 1
            cumulative_cf += cf
            discounted_cf = cf / ((1 + discount_rate) ** year)
            
            yearly_data.append({
                "year": year,
                "cash_flow": cf,
                "discounted_cash_flow": discounted_cf,
                "cumulative_cash_flow": cumulative_cf
            })
        
        # 返回结果
        return {
            "summary": {
                "initial_investment": initial_investment,
                "total_returns": total_returns,
                "net_profit": total_returns - initial_investment,
                "roi": roi,
                "npv": npv,
                "irr": irr if irr is not None else "N/A",
                "payback_period": payback_period if payback_period is not None else "N/A",
                "discount_rate": discount_rate,
                "project_years": project_years
            },
            "yearly_data": yearly_data,
            "charts": {
                "cash_flow": {
                    "type": "bar",
                    "data": cash_flow_values,
                    "labels": [f"Year {i+1}" for i in range(len(cash_flow_values))]
                },
                "cumulative_cash_flow": {
                    "type": "line",
                    "data": [d["cumulative_cash_flow"] for d in yearly_data],
                    "labels": [f"Year {i+1}" for i in range(len(yearly_data))]
                }
            }
        }
    
    def _calculate_cost_analysis(self, scene: Scene, data: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
        """
        成本测算计算逻辑
        
        Args:
            scene: 场景对象
            data: 输入数据
            
        Returns:
            计算结果数据
        """
        # 获取参数
        params = scene.parameters
        production_volume = params.get("production_volume", 1000)
        
        # 获取成本数据
        if "cost_data" not in data:
            raise ValueError("缺少成本数据")
        
        cost_data = data["cost_data"]
        
        # 计算固定成本和可变成本
        fixed_costs = cost_data[cost_data["cost_type"] == "fixed"]["amount"].sum()
        variable_cost_per_unit = cost_data[cost_data["cost_type"] == "variable"]["amount"].sum()
        
        # 计算总成本和单位成本
        total_variable_cost = variable_cost_per_unit * production_volume
        total_cost = fixed_costs + total_variable_cost
        unit_cost = total_cost / production_volume
        
        # 按成本类别汇总
        cost_by_category = cost_data.groupby("category").agg({"amount": "sum"}).reset_index()
        
        # 返回结果
        return {
            "summary": {
                "fixed_costs": fixed_costs,
                "variable_cost_per_unit": variable_cost_per_unit,
                "total_variable_cost": total_variable_cost,
                "total_cost": total_cost,
                "unit_cost": unit_cost,
                "production_volume": production_volume
            },
            "cost_breakdown": cost_by_category.to_dict(orient="records"),
            "charts": {
                "cost_by_category": {
                    "type": "pie",
                    "data": cost_by_category["amount"].tolist(),
                    "labels": cost_by_category["category"].tolist()
                },
                "cost_structure": {
                    "type": "bar",
                    "data": [fixed_costs, total_variable_cost],
                    "labels": ["Fixed Costs", "Variable Costs"]
                }
            }
        }
    
    def _calculate_profit_projection(self, scene: Scene, data: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
        """
        利润预测计算逻辑
        
        Args:
            scene: 场景对象
            data: 输入数据
            
        Returns:
            计算结果数据
        """
        # 获取参数
        params = scene.parameters
        forecast_periods = params.get("forecast_periods", 12)
        revenue_growth_rate = params.get("revenue_growth_rate", 0.05)
        cost_growth_rate = params.get("cost_growth_rate", 0.03)
        tax_rate = params.get("tax_rate", 0.25)
        
        # 获取收入和成本数据
        if "financial_data" not in data:
            raise ValueError("缺少财务数据")
        
        financial_data = data["financial_data"]
        
        # 获取最近一期的收入和成本
        last_revenue = financial_data["revenue"].iloc[-1]
        last_cost = financial_data["cost"].iloc[-1]
        
        # 创建预测期间
        forecast_dates = pd.date_range(
            start=pd.to_datetime(financial_data.index[-1]) + pd.DateOffset(months=1),
            periods=forecast_periods,
            freq='M'
        )
        
        # 预测收入和成本
        forecast_data = []
        for i in range(forecast_periods):
            # 计算预测收入和成本
            forecasted_revenue = last_revenue * (1 + revenue_growth_rate) ** (i + 1)
            forecasted_cost = last_cost * (1 + cost_growth_rate) ** (i + 1)
            
            # 计算毛利和税前利润
            gross_profit = forecasted_revenue - forecasted_cost
            profit_before_tax = gross_profit
            
            # 计算税金和净利润
            tax = profit_before_tax * tax_rate
            net_profit = profit_before_tax - tax
            
            # 计算利润率
            gross_margin = gross_profit / forecasted_revenue if forecasted_revenue > 0 else 0
            net_margin = net_profit / forecasted_revenue if forecasted_revenue > 0 else 0
            
            forecast_data.append({
                "date": forecast_dates[i].strftime("%Y-%m"),
                "revenue": forecasted_revenue,
                "cost": forecasted_cost,
                "gross_profit": gross_profit,
                "tax": tax,
                "net_profit": net_profit,
                "gross_margin": gross_margin,
                "net_margin": net_margin
            })
        
        # 计算汇总统计
        total_revenue = sum(item["revenue"] for item in forecast_data)
        total_cost = sum(item["cost"] for item in forecast_data)
        total_gross_profit = sum(item["gross_profit"] for item in forecast_data)
        total_net_profit = sum(item["net_profit"] for item in forecast_data)
        avg_gross_margin = total_gross_profit / total_revenue if total_revenue > 0 else 0
        avg_net_margin = total_net_profit / total_revenue if total_revenue > 0 else 0
        
        # 返回结果
        return {
            "forecast_data": forecast_data,
            "summary": {
                "total_revenue": total_revenue,
                "total_cost": total_cost,
                "total_gross_profit": total_gross_profit,
                "total_net_profit": total_net_profit,
                "avg_gross_margin": avg_gross_margin,
                "avg_net_margin": avg_net_margin,
                "forecast_periods": forecast_periods,
                "revenue_growth_rate": revenue_growth_rate,
                "cost_growth_rate": cost_growth_rate,
                "tax_rate": tax_rate
            },
            "charts": {
                "revenue_and_profit": {
                    "type": "line",
                    "data": {
                        "revenue": [item["revenue"] for item in forecast_data],
                        "gross_profit": [item["gross_profit"] for item in forecast_data],
                        "net_profit": [item["net_profit"] for item in forecast_data]
                    },
                    "labels": [item["date"] for item in forecast_data]
                },
                "profit_margins": {
                    "type": "line",
                    "data": {
                        "gross_margin": [item["gross_margin"] for item in forecast_data],
                        "net_margin": [item["net_margin"] for item in forecast_data]
                    },
                    "labels": [item["date"] for item in forecast_data]
                }
            }
        }
    
    def _calculate_custom(self, scene: Scene, data: Dict[str, pd.DataFrame]) -> Dict[str, Any]:
        """
        自定义测算逻辑
        
        Args:
            scene: 场景对象
            data: 输入数据
            
        Returns:
            计算结果数据
        """
        # 获取自定义计算逻辑
        params = scene.parameters
        custom_formula = params.get("custom_formula", "")
        
        if not custom_formula:
            raise ValueError("缺少自定义计算公式")
        
        # 创建局部变量空间，包含数据和常用函数
        local_vars = {
            "data": data,
            "np": np,
            "pd": pd
        }
        
        # 执行自定义计算逻辑
        try:
            # 使用exec执行自定义代码
            exec(custom_formula, {}, local_vars)
            
            # 检查是否生成了result变量
            if "result" not in local_vars:
                raise ValueError("自定义计算逻辑必须生成一个名为'result'的变量")
            
            result = local_vars["result"]
            
            # 确保结果是可序列化的
            if isinstance(result, pd.DataFrame):
                result = result.to_dict(orient="records")
            
            return {
                "custom_result": result,
                "formula": custom_formula
            }
            
        except Exception as e:
            raise ValueError(f"自定义计算逻辑执行错误: {str(e)}")